import arcpy
import os
import pollutantdata as pd

def CalculateLoad(inRasterPath, outRasterPath, pollutantConcentrationDictionary, rainfallDepth, gridCellArea, conversionFactor):
    landuseRaster = arcpy.sa.Raster(inRasterPath)

    print 'Picking pollutant concentrations'

    #This method relies on the landuse codes being mapped to the 'Value' field of the landuse raster such that
    # RES = 1, COM = 2, etc...
    #We could make this more generic but it's simple and it works for now.
    concentrationRaster = arcpy.sa.Pick(landuseRaster, [
        pollutantConcentrationDictionary['RES'].mean,
        pollutantConcentrationDictionary['COM'].mean,
        pollutantConcentrationDictionary['POS'].mean,
        pollutantConcentrationDictionary['IND'].mean,
        pollutantConcentrationDictionary['VAC'].mean,
        pollutantConcentrationDictionary['TRA'].mean,
        0]) #Landuse 'Value' of 7 corresponds to 'WAT', which always has a load of 0

    print 'Calculating pollutant load'
    loadRaster = concentrationRaster * rainfallDepth * gridCellArea * conversionFactor

    print 'Saving load raster: ' + outRasterPath
    loadRaster.save(outRasterPath)

    return

arcpy.CheckOutExtension("Spatial")
arcpy.env.overwriteOutput = True

#inRaster not the standard landuse raster - downsaved from 32bit for for faster processing
inRasterPath = r"\\ws18325\Stormwater D\RiskModel\Surface\LandClassification.gdb\LanduseZoningHybrid"
outputFolder = r"\\ws18325\Stormwater D\RiskModel\RiskAssessment\WaterQuality\LoadingEstimates.gdb"


CalculateLoad(inRasterPath, outputFolder + r'\tss_annual_load',
    pd.tssLuConcentration, pd.annualRainfallDepth, 9, pd.mg_per_L_conversion)

CalculateLoad(inRasterPath, outputFolder + r'\pb_annual_load',
    pd.pbLuConcentration, pd.annualRainfallDepth, 9, pd.ug_per_L_conversion)

CalculateLoad(inRasterPath, outputFolder + r'\tp_annual_load',
    pd.tpLuConcentration, pd.annualRainfallDepth, 9, pd.mg_per_L_conversion)

CalculateLoad(inRasterPath, outputFolder + r'\bod_annual_load',
    pd.bodLuConcentration, pd.annualRainfallDepth, 9, pd.mg_per_L_conversion)

CalculateLoad(inRasterPath, outputFolder + r'\ecoli_annual_load',
    pd.ecoliLuConcentration, pd.annualRainfallDepth, 9, pd.cfu_per_100_mL_conversion)
